import torch
from PIL import Image
from torchvision import transforms
import numpy as np

# ----------------------------
# 字符集配置
# ----------------------------
CHARS = list("0123456789一二三四五六七八九壹贰叁肆伍陆柒捌玖加减ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz家见等问")
BLANK_LABEL = 0
char2idx = {c: i + 1 for i, c in enumerate(CHARS)}
idx2char = {i + 1: c for i, c in enumerate(CHARS)}
idx2char[BLANK_LABEL] = ''

# ----------------------------
# 模型结构（保持与训练一致）
# ----------------------------
import torch.nn as nn

class CRNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2, 2),
        )
        self.rnn = nn.LSTM(64 * 10, 128, num_layers=2, bidirectional=True, batch_first=True)
        self.fc = nn.Linear(128 * 2, num_classes)

    def forward(self, x):
        b, c, h, w = x.size()
        x = self.cnn(x)
        x = x.permute(0, 3, 1, 2).contiguous()
        x = x.view(b, x.size(1), -1)
        x, _ = self.rnn(x)
        x = self.fc(x)
        return x.log_softmax(2)

# ----------------------------
# 解码函数：CTC greedy decode
# ----------------------------
def decode_output(logits):
    pred = logits.argmax(dim=2)[0].cpu().numpy()
    result = []
    prev = -1
    for p in pred:
        if p != prev and p != BLANK_LABEL:
            result.append(idx2char.get(p, ''))
        prev = p
    return ''.join(result)

# ----------------------------
# 图像预处理
# ----------------------------
transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((40, 100)),
    transforms.ToTensor(),
])

# ----------------------------
# 推理函数
# ----------------------------
def predict(image_path, model_path="captcha_crnn.pth"):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CRNN(len(CHARS) + 1)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval().to(device)

    image = Image.open(image_path).convert("L")
    image = transform(image).unsqueeze(0).to(device)  # [1, 1, 40, 100]

    with torch.no_grad():
        output = model(image)  # [1, T, C]
        pred_text = decode_output(output)
        print(f"✅ Predicted: {pred_text}")

# ----------------------------
# 执行测试
# ----------------------------
if __name__ == "__main__":
    predict("/temp/captcha/test/0332_0000088.png",model_path="captcha_crnn_best.pth")
